package com.yyds.tensordemo;

import android.content.Context;
import android.content.res.AssetManager;
import android.util.Log;

import com.google.gson.Gson;

import org.tensorflow.lite.DataType;
import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;

import java.io.FileDescriptor;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class InferenceHelper {
    private static final String TAG = "InferenceHelper";
    private Interpreter interpreter;

    /**
     * 加载 TensorFlow Lite 模型
     *
     * @param modelPath 模型文件路径
     */
    public void loadModel(String modelPath) {
        try {
            MappedByteBuffer tfliteModel = loadModelFile(modelPath);

            interpreter = new Interpreter(tfliteModel);

            Log.d(TAG, "Model loaded successfully.");
        } catch (IOException e) {
            Log.e(TAG, "Error loading model", e);
        }
    }

    public void loadModelFromAsset(Context context, String modelPath) {
        try {
            MappedByteBuffer tfliteModel = loadModelFileFromAsset(context, modelPath);

            interpreter = new Interpreter(tfliteModel);
            Log.d(TAG, "Model loaded successfully.");
        } catch (IOException e) {
            Log.e(TAG, "Error loading model", e);
        }
    }

    /**
     * 从文件路径加载模型
     *
     * @param modelPath 模型文件路径
     * @return MappedByteBuffer
     * @throws IOException
     */
    private MappedByteBuffer loadModelFile(String modelPath) throws IOException {
        FileChannel fileChannel = null;
        if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.O) {
            fileChannel = FileChannel.open(Paths.get(modelPath), StandardOpenOption.READ);
        }
        return fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size());
    }

    public static MappedByteBuffer loadModelFileFromAsset(Context context, String assetPath) throws IOException {
        // 获取AssetManager
        AssetManager assetManager = context.getAssets();

        // 打开Assets中的文件
        FileInputStream inputStream = assetManager.openFd(assetPath).createInputStream();

        // 获取FileDescriptor
        FileDescriptor fileDescriptor = inputStream.getFD();

        // 创建FileChannel
        FileChannel fileChannel = new FileInputStream(fileDescriptor).getChannel();

        // 映射整个文件到内存
        MappedByteBuffer mappedByteBuffer = fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size());

        // 关闭FileChannel和FileInputStream
        fileChannel.close();
        inputStream.close();

        return mappedByteBuffer;
    }

    /**
     * 使用模型进行推理
     *
     * @param features 输入的音频特征，形状为 [1, 100, 13]
     * @param mask     输入的 mask，形状为 [1, 100]
     * @return 预测的类别
     */

    String maskStr = "[[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]";
    String featuresStr = "[[[-398.6518859863281, 179.9567413330078, -25.54849624633789, 49.31053161621094, -16.591665267944336, 4.741461753845215, 9.852978706359863, 2.4031565189361572, -2.1452362537384033, -1.2605997323989868, -1.0094836950302124, 1.4413784742355347, 12.106977462768555], [-329.61468505859375, 175.503662109375, -31.851871490478516, 51.06398010253906, -14.567886352539062, 7.4163737297058105, 9.955795288085938, 5.058274269104004, -9.819658279418945, -2.1723241806030273, -4.015442371368408, 0.5346770286560059, 9.081537246704102], [-312.55780029296875, 172.22181701660156, -46.36773681640625, 49.402259826660156, -14.155593872070312, 5.670036792755127, 4.898420333862305, 13.587279319763184, -1.9523369073867798, -5.999358177185059, -11.70743465423584, -7.3457770347595215, -1.2621763944625854], [-277.4833984375, 191.9708251953125, -54.53293228149414, 54.03130340576172, -12.582164764404297, 2.0461301803588867, -2.993173599243164, 3.9835712909698486, -7.719878196716309, -7.786253929138184, -6.0989885330200195, -7.4419097900390625, -5.066836357116699], [-258.0925598144531, 198.53097534179688, -56.82081604003906, 52.92500305175781, -18.373933792114258, 3.7033603191375732, -1.9976956844329834, 4.239175796508789, -14.999863624572754, -12.691084861755371, -4.131605625152588, -6.184911727905273, -0.615018904209137], [-255.93553161621094, 192.80477905273438, -58.579627990722656, 54.79671096801758, -16.242259979248047, 7.450080394744873, -5.196530342102051, 4.395244598388672, -14.125914573669434, -16.22435760498047, -8.593301773071289, -9.820704460144043, -2.002445936203003], [-256.49542236328125, 185.9282684326172, -57.175010681152344, 58.26151657104492, -13.81134033203125, 6.773299217224121, -7.2660675048828125, -1.9652924537658691, -12.539968490600586, -20.825946807861328, -11.488468170166016, -15.307820320129395, -1.2768665552139282], [-230.11819458007812, 142.66551208496094, -19.940200805664062, 35.845970153808594, 9.610124588012695, -8.179767608642578, 0.6346516609191895, -4.695369720458984, -7.893675804138184, -16.72170639038086, -12.004898071289062, -11.541519165039062, -4.214258670806885], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]]";

    public int[] runInference(float[][][] features, float[][] mask) {

        Gson gson = new Gson();

        java.lang.reflect.Type type3 = new com.google.gson.reflect.TypeToken<List<List<List<Float>>>>() {
        }.getType();
        java.lang.reflect.Type type2 = new com.google.gson.reflect.TypeToken<List<List<Float>>>() {
        }.getType();
        List<List<List<Float>>> list3 = gson.fromJson(featuresStr, type3);
        List<List<Float>> list2 = gson.fromJson(maskStr, type2);

        features = new float[list3.size()][][];
        for (int i = 0; i < list3.size(); i++) {
            List<List<Float>> middleList = list3.get(i);
            features[i] = new float[middleList.size()][];
            for (int j = 0; j < middleList.size(); j++) {
                List<Float> innerList = middleList.get(j);
                features[i][j] = new float[innerList.size()];
                for (int k = 0; k < innerList.size(); k++) {
                    features[i][j][k] = innerList.get(k);
                }
            }
        }

        mask = new float[list2.size()][];
        for (int j = 0; j < list2.size(); j++) {
            List<Float> innerList = list2.get(j);
            mask[j] = new float[innerList.size()];
            for (int k = 0; k < innerList.size(); k++) {
                mask[j][k] = innerList.get(k);
            }
        }

// 展平三维数组
        TensorBuffer inputBufferF = TensorBuffer.createFixedSize(new int[]{1, 100, 13}, DataType.FLOAT32);
        ByteBuffer buffer = inputBufferF.getBuffer();
        float[][] array = features[0];
        for (int j = 0; j < array.length; j++) {
            for (int k = 0; k < array[j].length; k++) {
                buffer.putFloat(array[j][k]);
            }
        }
        buffer.position(0);
        Log.d("TAG", "Features shape: " + Arrays.toString(inputBufferF.getShape()));
// 展平二维数组
        TensorBuffer inputBufferMask = TensorBuffer.createFixedSize(new int[]{1, 100}, DataType.FLOAT32);
        ByteBuffer bufferMask = inputBufferMask.getBuffer();
        float[] arrayMask = mask[0];
        for (int j = 0; j < arrayMask.length; j++) {
            bufferMask.putFloat(arrayMask[j]);
        }
        bufferMask.position(0);
        Log.d("TAG", "Mask shape: " + Arrays.toString(inputBufferMask.getShape()));

        Map<Integer, Object> outputMap = new HashMap<>();
        int numOutputs = interpreter.getOutputTensorCount();
        for (int i = 0; i < numOutputs; i++) {
            TensorBuffer outputTensorBuffer = TensorBuffer.createFixedSize(interpreter.getOutputTensor(i).shape(), DataType.FLOAT32);
            outputMap.put(i, outputTensorBuffer.getBuffer()); // 将索引和缓冲区放入映射
            Log.d("TAG", "numOutputs shape: " + Arrays.toString(outputTensorBuffer.getShape()) + "   index:" + i);
        }
        Object[] inputArray = new Object[2];
        inputArray[0] = inputBufferF.getBuffer();
        inputArray[1] = inputBufferMask.getBuffer();
        // 输出提取的特征和掩码形状
        try {
            interpreter.runForMultipleInputsOutputs(inputArray, outputMap);
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    /**
     * 关闭推理器，释放资源
     */
    public void close() {
        if (interpreter != null) {
            interpreter.close();
            Log.d(TAG, "Interpreter closed.");
        }
    }
}